#include <stdio.h>
#include <errno.h>
#include <string.h>
#include <unistd.h>
#include <netdb.h>
#include <signal.h>
#include <stdlib.h>
#include <pthread.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <sys/time.h>
#include "my_icmp.h"
#include "my_ip.h"

int g_send_num; // 发包数
int g_recv_num; // 收包数

char g_host[64]; // 域名、主机名或者IP地址
struct sockaddr_in g_dst; // 目标主机

// 计算报文往返的RTT，单位毫秒
double round_trip_time(struct timeval *start, struct timeval *end)
{
	if (end->tv_usec < start->tv_usec) {
		end->tv_sec -= 1;
		end->tv_usec += 1000000;
	}
	return (end->tv_sec - start->tv_sec) * 1000 
			+ (end->tv_usec - start->tv_usec) / 1000;
}

/********************************************************
ICMP校验和算法：
1. 把校验和字段置为0
2. 对ICMP头部中的每16bit进行求和
3. 如果和的高16bit不为0，则将和的高16bit和低16bit反复相加，
   直到和的高16bit为0
4. 将该16bit的值取反，存入校验和字段
*********************************************************/
unsigned short checksum(unsigned short *addr, int size)
{
	unsigned int sum = 0;
	
	// 把ICMP报头二进制数据以2字节为单位累加起来
	while (size > 1) {
		sum += *addr++;
		size -= 2;
	}

	// 如果ICMP报头为奇数个字节，会剩下最后1字节
	if (size > 0) {
		// 把最后1个字节视为一个2字节数据的高字节，低字节为0，累加
		unsigned short left = 0;
		*(unsigned char *)(&left) = *(unsigned char *)addr;
		sum += left;
	}

	while (sum >> 16) {
		// 将高16bit与低16bit相加
		sum = (sum >> 16) + (sum & 0xffff);
	}

	return ~sum;
}

int pack(int seq, icmp_hdr *icmp)
{
	icmp->type = ICMP_ECHO_REQUEST;
	icmp->code = 0;
	icmp->checksum = 0;
	icmp->id = getpid();
	icmp->seq = seq;

	struct timeval start;
	gettimeofday(&start, NULL);
	memcpy(icmp->data, &start, sizeof(start));
	
	int size = sizeof(*icmp) + sizeof(start);
	icmp->checksum = checksum((unsigned short *)icmp, size);
	return size;
}

int unpack(char *buf, int len)
{
	ip_hdr *ip = (ip_hdr *)buf;
	int ip_hdr_len = ip->hl << 2;
	icmp_hdr *icmp = (icmp_hdr *)(buf + ip_hdr_len);

	len -= ip_hdr_len;
	if (len < 8) {
		printf("ICMP packet length error\n");
		return -1;
	}

	// 处理ICMP应答消息，只处理本进程发出消息的应答消息
	if (icmp->type == ICMP_ECHO_REPLY && icmp->id == getpid()) {
		struct timeval end;
		gettimeofday(&end, NULL);
		
		struct timeval start;
		memcpy(&start, icmp->data, sizeof(start));

		// 计算PING报文的RTT
		double rtt = round_trip_time(&start, &end);
		
		printf("%d bytes from %s: icmp_seq=%u ttl=%d time=%.1f ms\n",
			len, inet_ntoa(g_dst.sin_addr), icmp->seq, ip->ttl, rtt);
	}

	return 0;
}

void *send_request(void *arg)
{
	int ret;
	int size;
	char packet[1024];
	int sockfd = *(int *)arg;

	for (int i = 0; i < 4; i++) {
		g_send_num++;
		memset(packet, 0, sizeof(packet));
		size = pack(i + 1, (icmp_hdr *)packet);
		ret = sendto(sockfd, packet, size, 0, (struct sockaddr *)&g_dst, sizeof(g_dst));
		if (ret < 0) {
			perror("sendto error");
			continue;
		}
		sleep(1);
	}

	return NULL;
}

void *recv_reply(void *arg)
{
	int ret;
	char packet[1024];
	int sockfd = *(int *)arg;
	struct sockaddr_in from;
	socklen_t len = sizeof(from);

	for (int i = 0; i < 4; i++) {
		ret = recvfrom(sockfd, packet, sizeof(packet), 0, (struct sockaddr *)&from, &len);
		if (ret < 0) {
			if (errno == EAGAIN) {
				continue;
			} else {
				perror("recvfrom error");
				continue;
			}
		}
		
		if (unpack(packet, ret) == -1) {
			printf("unpack error\n");
			continue;
		}
		g_recv_num++;
	}

	return NULL;
}

void statistic()
{
	printf("\n--- %s ping statistics ---\n", g_host);
	printf("%d packets transmitted, %d received, %%%d packet loss\n", 
		g_send_num, g_recv_num, (g_send_num - g_recv_num) * 100 / g_send_num);
}

void sig_handler(int signo)
{
	statistic();
	exit(1);
}

int main(int argc, char *argv[])
{
	if (argc < 2) {
		printf("invalid parameter\n");
		return -1;
	}
	memcpy(g_host, argv[1], strlen(argv[1]));

	// 通过信号结束进程，输出收发包统计信息
	signal(SIGINT, sig_handler);
	signal(SIGQUIT, sig_handler);
	signal(SIGTERM, sig_handler);

	// 创建原始套接字，协议类型是ICMP
	int sockfd = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP);
	if (sockfd < 0) {
		perror("socket error");
		return -1;
	}

	// 收包超时时间设置为23秒
	struct timeval timeout = {5, 0};
	setsockopt(sockfd, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));

	g_dst.sin_family = AF_INET;
	g_dst.sin_addr.s_addr = inet_addr(argv[1]);
	if (g_dst.sin_addr.s_addr == -1) {
		// 通过域名获取IP地址
		struct hostent *host = gethostbyname(argv[1]);
		if (host == NULL) {
			printf("invalid host name\n");
			close(sockfd);
			return -1;
		}
		memcpy(&g_dst.sin_addr, host->h_addr, host->h_length);
	}

	// 发包线程
	pthread_t stid;
	pthread_create(&stid, NULL, send_request, &sockfd);

	// 收包线程
	pthread_t rtid;
	pthread_create(&rtid, NULL, recv_reply, &sockfd);

	pthread_join(stid, NULL);
	pthread_join(rtid, NULL);

	// 收发报文统计信息
	statistic();
	close(sockfd);
	return 0;
}
